package task_pool

import (
	"sync"
)

// TaskFunc 需要被并发执行的函数
type TaskFunc[T any] func() (T, error)

// TaskResult 封装任务执行结果
type TaskResult[T any] struct {
	Value T
	Err   error
}

// TaskPool 线程池
type TaskPool[T any] struct {
	taskChan     chan TaskFunc[T]   // 任务管道，管道中是待执行的数据
	resChan      chan TaskResult[T] // 返回值管道，管道中存储所有任务的返回值
	wg           *sync.WaitGroup    // 并发控制器
	coroutineNum int                // 最大的并发协程数量
	taskNum      int                // 任务数量
}

// NewTaskPool 新建一个线程池对象，并开始运行
func NewTaskPool[T any](coroutineNum, taskNum int) *TaskPool[T] {
	pool := &TaskPool[T]{
		taskChan:     make(chan TaskFunc[T], taskNum),
		resChan:      make(chan TaskResult[T], taskNum),
		wg:           &sync.WaitGroup{},
		coroutineNum: coroutineNum,
		taskNum:      taskNum,
	}
	pool.wg.Add(pool.coroutineNum)

	return pool
}

// Run 开始运行任务
func (pool *TaskPool[T]) Run() {
	for i := 0; i < pool.coroutineNum; i++ {
		go func() {
			defer pool.wg.Done()
			for task := range pool.taskChan {
				val, err := task()
				pool.resChan <- TaskResult[T]{Value: val, Err: err}
			}
		}()
	}
}

// SubmitAllTask 向线程池提交所有需要被执行的任务
func (pool *TaskPool[T]) SubmitAllTask(tasks []TaskFunc[T]) {
	// 将所有任务提交至 taskChan 管道
	for _, task := range tasks {
		pool.taskChan <- task
	}
	// 关闭任务管道，防止冲突
	close(pool.taskChan)
}

// Close 关闭线程池
func (pool *TaskPool[T]) Close() {
	pool.wg.Wait()
	close(pool.resChan)
}
